import json
import os
from typing import Type, Any

from langchain import requests
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import BaseTool

from internal.lib import decorate


class CustomModel(BaseModel):
    city: str = Field(description="要查询的城市的名称")


class CustomTool(BaseTool):
    name: str = "weather_info",
    description: str = "根据城市查询天气预报信息",
    args_schema: Type[BaseModel] = CustomModel
    return_direct: bool = True

    def _run(self, *args: Any, **kwargs: Any) -> str:
        city = kwargs.get("city")
        url = f"https://restapi.amap.com/v3"
        key = os.getenv("GAODE_API_KEY")
        if key is None:
            return "密钥不存在或者错误"
        session = requests.session()
        try:
            ad_code_obj = session.request(
                url=f"{url}/config/district?key={key}&keywords={city}&subdistrict=0&extensions=all",
                method="GET",
                headers={"Content-Type": "application/json; charset=utf-8"}
            )
            ad_code_obj.raise_for_status()
            ad_code_json = ad_code_obj.json()
            if ad_code_json.get("info") == "OK":
                adcode = ad_code_json["districts"][0].get("adcode")

                weather = session.request(
                    url=f"{url}/weather/weatherInfo?key={key}&city={adcode}&extensions=all",
                    method="GET",
                    headers={"Content-Type": "application/json; charset=utf-8"}
                )
                weather.raise_for_status()
                weather_json = weather.json()
                if weather_json.get("info") == "OK":
                    return json.dumps(weather_json)
                return f"查询{city}天气预报失败"
        except Exception as e:
            return f"查询{city}天气预报失败"


@decorate("args_schema", CustomModel)
def gaode_weather(**kwargs) -> BaseTool:
    return CustomTool()
